{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a96f3ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "#lightning\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import seaborn as sns\n",
    "class dla(object):\n",
    "    def __init__(self,size=[100,100],core_pos=[[50,50]]):\n",
    "        self.map=np.zeros(size)\n",
    "        for c in core_pos:\n",
    "            #print(c)\n",
    "            self.map[c[0],c[1]]=1\n",
    "    def main(self,time=100,plt_show=True):\n",
    "        for t in range(time):\n",
    "            self.random_start()\n",
    "            print(\"No.\"+str(t)+\" is running...\")\n",
    "            rota=True\n",
    "            while rota:\n",
    "                self.random_walk()\n",
    "                if self.near(self.point):\n",
    "                    self.map[self.point[0],self.point[1]]=1\n",
    "                    rota=False\n",
    "            if plt_show:\n",
    "                self.map_show(i=t)\n",
    "        return None\n",
    "    def random_walk(self):\n",
    "        move=np.random.randint(4)\n",
    "        if move==0 and self.point[0]-1>-1:\n",
    "            self.map[self.point[0],self.point[1]]=0\n",
    "            self.point[0]-=1\n",
    "            self.map[self.point[0],self.point[1]]=0.5\n",
    "        if move==1 and self.point[0]+1<len(self.map):\n",
    "            self.map[self.point[0],self.point[1]]=0\n",
    "            self.point[0]+=1\n",
    "            self.map[self.point[0],self.point[1]]=0.5\n",
    "        if move==2 and self.point[1]-1>-1:\n",
    "            self.map[self.point[0],self.point[1]]=0\n",
    "            self.point[1]-=1\n",
    "            self.map[self.point[0],self.point[1]]=0.5\n",
    "        if move==3 and self.point[1]+1<len(self.map[0]):\n",
    "            self.map[self.point[0],self.point[1]]=0\n",
    "            self.point[1]+=1\n",
    "            self.map[self.point[0],self.point[1]]=0.5\n",
    "        return None\n",
    "    def random_start(self):\n",
    "        while True:\n",
    "            self.point=[np.random.randint(len(self.map)),np.random.randint(len(self.map[0]))]\n",
    "            if self.map[self.point[0],self.point[1]]==0 and not self.near(self.point):\n",
    "                self.map[self.point[0],self.point[1]]=0.5\n",
    "                break\n",
    "        return None             \n",
    "    def near(self,pos):\n",
    "        if pos[0]+1<len(self.map) and pos[1]+1<len(self.map[0]) and pos[0]-1>-1 and pos[1]-1>-1:\n",
    "            if self.map[pos[0]+1,pos[1]]==1 or self.map[pos[0]-1,pos[1]]==1 or self.map[pos[0],pos[1]+1]==1 or self.map[pos[0],pos[1]-1]==1:\n",
    "                return True\n",
    "            else:\n",
    "                return False\n",
    "        else:\n",
    "            return False\n",
    "    def map_show(self,i):\n",
    "        fig = plt.figure(figsize=(12,10))\n",
    "        ax = sns.heatmap(self.map,cmap=\"binary\") \n",
    "        plt.savefig(\"D:/3_bodies/DLA_pic/single_\"+str(i)+\".jpg\")\n",
    "        plt.show()\n",
    "def gen_core(core_img,size=[500,500]):\n",
    "    core_pos=[]\n",
    "    for i in range(size[0]):\n",
    "        core_pos.append([i,0])\n",
    "        core_pos.append([i,size[0]-1])\n",
    "    for i in range(size[1]):\n",
    "        core_pos.append([size[0]-1,i])\n",
    "    for i in range(len(core_img)):\n",
    "        for j in range(len(core_img[0])):\n",
    "            if core_img[i,j].all()==0:\n",
    "                core_pos.append([-int((size[0]+len(core_img))/2-i),int((size[1]-len(core_img[0]))/2+j)])\n",
    "    return core_pos\n",
    "def img2bi(img_path,size=[500,500],save_path=None):\n",
    "    img=Image.open(img_path)\n",
    "    img=img.resize(size)\n",
    "    arr=np.asarray(img)\n",
    "    ave=np.mean(arr)\n",
    "    print(arr.shape)\n",
    "    arr_bi=np.zeros(arr.shape)\n",
    "    for i in range(len(arr)):\n",
    "        for j in range(len(arr[i])):\n",
    "            if np.mean(arr[i,j])>ave:\n",
    "                arr_bi[i,j]=1\n",
    "            else:\n",
    "                arr_bi[i,j]=0\n",
    "    if save_path!=None:\n",
    "        arr_bi*=255\n",
    "        image=Image.fromarray(arr_bi.astype('uint8')).convert('RGB')\n",
    "        image.save(save_path)\n",
    "        image.show()\n",
    "    return arr_bi\n",
    "img_path=\"D:\\\\3_bodies\\\\DLA\\\\heart.png\"\n",
    "g_arr=img2bi(img_path=img_path,size=[300,300])#\n",
    "core_pos=gen_core(g_arr)\n",
    "#print(core_pos)\n",
    "d_0=dla(size=[100,100],core_pos=[50,50])#core_pos=[[25,25],[75,25],[25,75],[75,75]]\n",
    "d_0.main(plt_show=True,time=10000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
